import numpy as np
import pandas as pd
import os
from kulfan_to_coord import CST_shape
from scipy.optimize import leastsq, least_squares
import matplotlib.pyplot as plt
import shutil

var_nums = 10


class CST:
    def __init__(self, dataset_dir, v_save_dir, var_nums):
        self.dataset_dir = dataset_dir
        self.v_save_dir = v_save_dir
        self.var_nums = var_nums
        self.dz = 0
        self.N = None
        self.x1 = None
        self.y1 = None
        self.plsq = None
        self.airfoil_CST = None

    def read_data(self, file_path):
        ordinates = pd.read_csv(file_path).to_numpy()
        self.x1 = ordinates[:, 0]
        self.y1 = ordinates[:, 1]
        self.N = ordinates.shape[0]

    def _objfunc(self, var):
        var_nums = int(self.var_nums)
        wl = var[0:var_nums]
        wu = var[var_nums:2 * var_nums]
        wl[0] = abs(wl[0])
        wu[0] = -wl[0]

        self.airfoil_CST = CST_shape(wl, wu, self.dz, self.N)
        coordinates = self.airfoil_CST.inv_airfoil_coor(self.x1)
        y2 = coordinates[:][1]

        return y2 - self.y1

    def cal_error(self):
        dy = np.abs(self._objfunc(self.plsq))  ###
        range1 = np.where(self.x1 < 0.2)[0]
        range2 = np.where(self.x1 >= 0.2)[0]
        max1 = np.max(dy[range1])
        max2 = np.max(dy[range2])

        print("x < 0.2:", max1 < 7e-4)
        print("x >= 0.2:", max2 < 4e-4)
        print("x < 0.2 (max):", max1)
        print("x >= 0.2 (max):", max2)
        return max1 < 7e-4, max2 < 4e-4, max1, max2

    def cal_mean_error(self):
        dy = np.abs(self._objfunc(self.plsq))  ###
        range1 = np.where(self.x1 < 0.2)[0]
        range2 = np.where(self.x1 >= 0.2)[0]
        mean1 = np.mean(dy[range1])
        mean2 = np.mean(dy[range2])

        print("x < 0.2 (mean):", mean1)
        print("x >= 0.2 (mean):", mean2)
        return mean1, mean2

    def analyse(self):
        dy = np.abs(self._objfunc(self.plsq))
        range1 = np.where(self.x1 < 0.2)[0]
        range2 = np.where(self.x1 >= 0.2)[0]
        self.error1_id = np.intersect1d(np.where(dy >= 7e-4), range1)
        self.error2_id = np.intersect1d(np.where(dy >= 4e-4), range2)
        # print(self.error1_id)
        # print(self.error2_id)

    def analyse_max(self):

        dy = np.abs(self._objfunc(self.plsq))
        range1 = np.where(self.x1 < 0.2)[0]
        range2 = np.where(self.x1 >= 0.2)[0]
        # error1 = np.where(dy[range1] >= 7e-4)
        # error2 = np.where(dy[range2] >= 4e-4)
        error1_max_id = range1[np.argmax(dy[range1])]
        error2_max_id = range2[np.argmax(dy[range2])]

        return [self.x1[error1_max_id], self.x1[error2_max_id]]

    def optimize(self):
        p_init = np.random.randn(2 * self.var_nums)  # 初始化
        self.plsq = leastsq(func=self._objfunc, x0=p_init, ftol=1.49012e-18, xtol=1.49012e-18)[0]

    def optimize_v2(self):
        p_init = np.random.randn(2 * self.var_nums)  # 初始化
        self.plsq = least_squares(self._objfunc,
                                  p_init,
                                  verbose=0,
                                  ftol=1e-15,
                                  xtol=1e-15,
                                  gtol=1e-15,
                                  loss="arctan",
                                  max_nfev=100,
                                  diff_step=100,
                                  method="trf").x

    def get_w(self):
        wl_new = self.plsq[0:self.var_nums]
        wu_new = self.plsq[self.var_nums:2 * self.var_nums]
        return wl_new, wu_new

    def loop_dataset(self):
        save_error_path1 = r"error/err1_{}.csv".format(self.var_nums)
        save_error_path2 = r"error/err2_{}.csv".format(self.var_nums)
        df1 = pd.DataFrame()
        df2 = pd.DataFrame()
        df = pd.DataFrame()  # columns=["file_name"] + ["x{}".format(i) for i in range(1, self.var_nums*2+1)])
        count = 0
        error_num_threhold = 5
        allow_error_num = 0
        error_max_x_list = []
        error_mean1_list = []
        error_mean2_list = []

        save_path = os.path.join(self.v_save_dir, "vars_{}.csv".format(self.var_nums))
        for file in os.listdir(self.dataset_dir):
            # for file in ["goe701.csv"]:
            file_path = os.path.join(self.dataset_dir, file)
            print(file_path + ":")

            self.read_data(file_path)
            self.optimize_v2()  # optimize
            tolerance1, tolerance2, max1, max2 = self.cal_error()

            # 计算不符合误差的个数
            # 将满足误差和不满足误差的翼型分开
            if not tolerance2 or not tolerance1:
                count = count + 1
                shutil.copy(file_path, os.path.join("beyondError", file))
            else:
                shutil.copy(file_path, os.path.join("underError", file))

                # 设计变量
                wl_new, lu_new = obj.get_w()
                w_new = np.concatenate((wl_new, lu_new), axis=0).T
                df = df.append([[file] + list(w_new)])

            # 计算平均误差
            mean1, mean2 = self.cal_mean_error()
            error_mean1_list.append(mean1)
            error_mean2_list.append(mean2)

            # 每个翼型是否满足误差
            df1 = df1.append([[file, tolerance1, tolerance2]])

            # 每个翼型的最大误差
            df2 = df2.append([[file, max1, max2]])

            # 计算哪些点超出设定的阈值
            self.analyse()
            self.cst_save_plot(file)

            # 计算有几个点超过误差了
            if self.error1_id.shape[0] + self.error2_id.shape[0] <= error_num_threhold:
                allow_error_num += 1
            error_max_x_list.append(self.analyse_max())
            break

        error_max_x_list = np.asarray(error_max_x_list)
        print("save var to:", save_path)
        df.to_csv(save_path, index=False, header=True)  # 保存设计变量
        # df1.to_csv(save_error_path1, index=0)  # x<0.2误差
        # df2.to_csv(save_error_path2, index=0)  # x>=0.2误差

        print("不符合限定误差的样本个数:", count)
        print("误差点个数在{}个以内的样本数".format(error_num_threhold), allow_error_num)
        print("整个数据集平均误差：x<=0.2:{};x>=0.2:{}".format(np.mean(error_mean1_list),
                                                     np.mean(error_mean2_list)))

        # print("误差最大值的x坐标分布：")
        # plt.hist(error_max_x_list[:, 0], bins=10)
        # plt.hist(error_max_x_list[:, 1], bins=40)
        # plt.show()
        # plt.cla()

        # out_of_range_error1 = error_max_x_list[np.where(error_max_x_list[:, 0] >= 7e-4)[0], 0]
        # out_of_range_error2 = error_max_x_list[np.where(error_max_x_list[:, 1] >= 4e-4)[0], 1]
        # plt.hist(out_of_range_error1, bins=10)
        # plt.hist(out_of_range_error2, bins=40)
        # plt.show()

    def cst_save_plot(self, file):

        wl_new = self.plsq[0:self.var_nums]
        wu_new = self.plsq[self.var_nums:2 * self.var_nums]
        for i in self.plsq:
            print("{},".format(i), end="")

        airfoil_CST2 = CST_shape(wl_new, wu_new, 0, N)
        coordinates = airfoil_CST2.airfoil_coor()
        print(coordinates)

        x_coor = coordinates[0]
        y_coor = coordinates[1]
        np.set_printoptions(suppress=True)
        _, _, tmp1, tmp2 = self.cal_error()
        file = "{0:.3}_{1:.3}_{2}".format(np.mean(tmp1),
                                          np.mean(tmp2),
                                          file)

        # dat file
        output_dir = r"CST_output"
        file_path = os.path.join(output_dir, file)
        pd.DataFrame({"x": np.asarray(x_coor).squeeze(), "y": np.asarray(y_coor).squeeze()}).to_csv(
            file_path, index=False, header=False)

        # plot
        # plt.title(file)
        # ax = plt.subplot(111)
        # ax.plot(x_coor, y_coor, 'g-', linewidth=4, label='CST')
        # ax.plot(self.x1, self.y1, 'r', label='original')
        # if self.error1_id is not None and self.error2_id is not None:
        #     ax.scatter(self.x1[self.error1_id], self.y1[self.error1_id], c='b')
        #     ax.scatter(self.x1[self.error2_id], self.y1[self.error2_id], c='r')
        #
        # ax.legend(loc='lower center', frameon=False)
        # plt.xlabel('x/c')
        # plt.ylabel('y/c')
        # plt.xlim((0, 1.0))  # 设置横坐标范围
        # plt.ylim((-0.3, 0.3))  # 设置纵坐标范围
        # plt.xticks(np.arange(0, 1, 0.1))
        # plt.gca().set_aspect(1)
        # plt.grid()  # 设置网格
        # ax.spines['right'].set_visible(False)
        # ax.spines['top'].set_visible(False)
        # ax.yaxis.set_ticks_position('left')
        # ax.xaxis.set_ticks_position('bottom')
        # plt.savefig(os.path.join(r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_6_ssp4_show", file[:-3] + "jpg"))
        # plt.clf()
        # plt.show()

    def curve2D_CST(self, x, paraVar, n1, n2, tolerance=1.0E-8):
        '''
        CST functions.

        x: X coordinate of curve, 0<= x <=1. 1Darray.
        paraVar: parameterization variables. 1Darray，只有大小没有位置
                 example:
                 array([y1,y2,...,yn])
        n1,n2: 整型变量，控制类型函数.
        还需要什么变量自己定义

        return
        -----------------------------------------
        dyMatrix: 2Darray. len(paraVar)行，len(x)列。
                  array([
                        [用第一组参数化变量生成的曲线],
                        [用第二组参数化变量生成的曲线],
                        ...,
                        []
                        ])
        dy: dyMatrix的列求和。1Darray，len(x)个元素。
        '''
        nPara = len(paraVar)
        # 仿照下面的式子写，应该一个式子能搞定

        # 列求和
        dy = dyMatrix.sum(axis=0)

        # abs<=tolerance, equal zero
        dy[np.abs(dy) < tolerance] = 0.0
        dyMatrix[np.abs(dyMatrix) < tolerance] = 0.0

        return dy, dyMatrix


if __name__ == "__main__":
    dataset_dir = r"D:\data_cmp\airfoil_csv_op\NACA4_ssp"  # r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_6_ssp4"
    v_save_dir = r"vars"
    N = 301
    obj = CST(dataset_dir=dataset_dir,
              v_save_dir=v_save_dir,
              var_nums=var_nums)
    obj.loop_dataset()
